package com.gitee.wsl.ext.collection

import com.gitee.wsl.data.cursor.cursor
import kotlin.reflect.KClass
import kotlin.reflect.cast
import kotlin.sequences.flatMap

fun <T, U, R> Sequence<T>.product(
    other: Sequence<U>,
    combinator: (T, U) -> R,
): Sequence<R> =
    flatMap { x ->
        other.map { y -> combinator(x, y) }
    }

fun <T, U> Sequence<T>.product(other: Sequence<U>): Sequence<Pair<T, U>> = product(other, ::Pair)

fun <T, R> Sequence<T>.squared(combinator: (T, T) -> R): Sequence<R> = product(this, combinator)

fun <T> Sequence<T>.squared(): Sequence<Pair<T, T>> = product(this)

/*fun <T> Sequence<T>.longIndexed(): Sequence<Pair<Long,T>> =
    zip(LongRange(0, Long.MAX_VALUE).asSequence()) { it, i ->
        LongIndexed.of(i, it)
    }*/

fun <T> Sequence<T>.indexed(): Sequence<Pair<Int,T>> =
    zip(IntRange(0, Int.MAX_VALUE).asSequence()) { it, i ->
        i to it
    }

fun <T> Sequence<T>.subsequences(): Sequence<Sequence<T>> {
    return sequence {
        var maxSize = 1
        var actualSize = 0
        while (true) {
            val sublist = this@subsequences.take(maxSize).toList()
            yield(sublist.asSequence())
            if (actualSize >= sublist.size) {
                break
            }
            maxSize++
            actualSize = sublist.size
        }
    }
}

fun <T> Sequence<T>.buffered(): Sequence<T> {
    return this.toList().asSequence()
}

fun <T> Sequence<T>.cached(): Sequence<T> {
    return this.cursor().asSequence()
}

fun <T> Sequence<T>.skipIndex(index: Int): Sequence<T> {
    require(index >= 0)
    return sequence {
        var i = 0
        val iter = iterator()
        while (i < index && iter.hasNext()) {
            yield(iter.next())
            i++
        }
        if (iter.hasNext()) iter.next()
        while (iter.hasNext()) {
            yield(iter.next())
        }
    }
}

fun <T> Sequence<T>.insertAt(
    index: Int,
    item: T,
    vararg items: T,
): Sequence<T> =
    sequence {
        for ((i, x) in withIndex().asIterable()) {
            if (i == index) {
                yield(item)
                yieldAll(items.asIterable())
            }
            yield(x)
        }
    }

fun <T> Sequence<T>.dropLast(): Sequence<T> =
    sequence {
        val i = iterator()
        while (i.hasNext()) {
            val current = i.next()
            if (i.hasNext()) {
                yield(current)
            }
        }
    }


/**
 * Returns a sequence that ignores first [count] elements.
 * Throws [IllegalArgumentException] if [count] is negative.
 */
 fun <T> Sequence<T>.drop(count: Long): Sequence<T> {
    require(count >= 0) { "Drop count should be non-negative, but had $count" }
    return sequence {
        var skipped = 0L
        forEach { value ->
            if (skipped >= count) yield(value) else ++skipped
        }
    }
}

fun <T> Sequence<T?>.assertItemsAreNotNull(): Sequence<T> =
    indexed().map {
        when (val value = it.second) {
            null -> throw IllegalArgumentException("Item at index ${it.first} is null")
            else -> value
        }
    }

fun <T> Sequence<T>.peek(action: (T) -> Unit) = map { action(it); it }

fun <T> Sequence<T>.toIndexedString(prefix: String = ""): String =
    mapIndexed { index, recognizerMatch -> index to recognizerMatch }
        .joinToString(separator = "\n", prefix = prefix) { "${it.first}: ${it.second}" }

 fun <T : Comparable<T>> Sequence<T>.isSorted(): Boolean =
    isSortedBy { it }

inline fun <T, R : Comparable<R>> Sequence<T>.isSortedBy(crossinline selector: (T) -> R?): Boolean =
    isSortedWith(compareBy(selector))

fun <T> Sequence<T>.isSortedWith(comparator: Comparator<in T>): Boolean =
    zipWithNext { a, b -> comparator.compare(a, b) <= 0 }.all { it }

fun <T : Comparable<T>> Sequence<T>.isSortedDescending(): Boolean =
    isSortedWith(reverseOrder())

inline fun <T, R : Comparable<R>> Sequence<T>.isSortedByDescending(crossinline selector: (T) -> R?): Boolean =
    isSortedWith(compareByDescending(selector))

inline fun <T> Sequence<T>.padWith(size: Int, inserter: (Sequence<T>) -> Sequence<T>): Sequence<T> {
    var result = this
    while (result.count() < size) {
        result = inserter(result)
    }
    return result
}

inline fun <T> Sequence<T>.padEnd(size: Int, padBlock: (Int) -> T): Sequence<T> = padWith(size) { it + padBlock(it.count()) }

 fun <T> Sequence<T>.padEnd(size: Int, o: T) = padEnd(size) { o }

 fun <T> Sequence<T>.padStart(size: Int, padBlock: (Int) -> T): Sequence<T> = padWith(size) { sequenceOf(padBlock(it.count())) + it }

 fun <T> Sequence<T>.padStart(size: Int, o: T) = padStart(size) { o }

fun <C : Comparable<C>> Sequence<C?>.indexOfMax(): Int {
    val iterator = iterator()
    if (!iterator.hasNext()) return -1
    var value = iterator.next()
    var index = 0
    while (value == null) {
        if (!iterator.hasNext()) return -1
        value = iterator.next()
        index++
    }
    var max: C = value
    var maxIndex = index
    if (!iterator.hasNext()) return maxIndex
    do {
        val v = iterator.next()
        index++
        if (v != null && max < v) {
            max = v
            maxIndex = index
        }
    } while (iterator.hasNext())
    return maxIndex
}


/**
 * Returns first element that is instance of specific type parameter [R], or `null` if not found
 */
inline fun <reified R> Sequence<*>.firstInstance(): R? = filterIsInstance<R>().firstOrNull()

inline fun <K, V, reified R> Sequence<Map.Entry<K, V>>.filterValuesIsInstance(): Sequence<Map.Entry<K, R>> {
    return filter { it.value is R } as Sequence<Map.Entry<K, R>>
}

inline fun <reified T : Any> Sequence<*>.firstIsInstanceOrNull(): T? {
    for (element in this) if (element is T) return element
    return null
}

fun <T, K> Sequence<T>.findDuplicates(by: (T) -> K): MutableMap<K, HashSet<T>> {
    val seen = mutableSetOf<K>()
    val duplicates = mutableMapOf<K, HashSet<T>>()

    for (item in this) {
        val key = by(item)
        if (!seen.add(key)) { // If `add` returns false, it's a duplicate
            val listForKey = duplicates.getOrElse(key) { hashSetOf() }
            listForKey.add(item)
            duplicates[key] = listForKey
        }
    }

    return duplicates
}


 fun <T> Sequence<T>.take(n: Int, default: T): Sequence<T> = sequence {
    val result = take(n).toList()
    yieldAll(result)
    yieldAll(List(n - result.size) { default })
}

 fun <T> Sequence<T>.isDistinct(): Boolean =
    toList().isDistinct()

 inline fun <T, K> Sequence<T>.isDistinctBy(crossinline keySelector: (T) -> K): Boolean =
    map { keySelector(it) }.isDistinct()

 inline fun <T, K : Any> Sequence<T>.groupByNotNull(crossinline keySelector: (T) -> K?): Map<K, List<T>> =
    mapNotNull { (keySelector(it) ?: return@mapNotNull null) to it }
        .groupBy({ it.first }, { it.second })

inline fun <T, K : Any, V : Any> Sequence<T>.groupByNotNull(
    crossinline keySelector: (T) -> K?,
    crossinline valueTransform: (T) -> V?
): Map<K, List<V>> =
    mapNotNull {
        val k = keySelector(it) ?: return@mapNotNull null
        val v = valueTransform(it) ?: return@mapNotNull null
        k to v
    }.groupBy({ it.first }, { it.second })

 inline fun <T, K : Any, V : Any> Sequence<T>.associateNotNull(
    crossinline valueSelector: (T) -> Pair<K?, V?>?,
): Map<K, V> = mapNotNull {
    val (k, v) = valueSelector(it) ?: return@mapNotNull null
    if (k == null || v == null) return@mapNotNull null
    k to v
}.toMap()

 inline fun <T, K : Any> Sequence<T>.associateByNotNull(
    crossinline keySelector: (T) -> K?,
): Map<K, T> = mapNotNull {
    (keySelector(it) ?: return@mapNotNull null) to it
}.toMap()

 inline fun <T, V : Any> Sequence<T>.associateWithNotNull(
    crossinline valueSelector: (T) -> V?,
): Map<T, V> = mapNotNull {
    it to (valueSelector(it) ?: return@mapNotNull null)
}.toMap()

fun <T> Sequence<T>.chunkedUntil(predicate: (T) -> Boolean): Sequence<List<T>> {
    return sequence {
        var list = mutableListOf<T>()
        this@chunkedUntil.forEach {
            if (!predicate(it)) {
                list.add(it)
            } else {
                yield(list)
                list = mutableListOf(it)
            }
        }
        if (list.isNotEmpty()) yield(list)
    }
}

fun <T> Sequence<T>.sumOf(selector: (T) -> Float): Float {
    var sum = 0f
    for (element in this) {
        sum += selector(element)
    }
    return sum
}

fun <T, R> Sequence<T>.flatMapTranArray(transform: (T) -> Array<R>): Sequence<R> =
    flatMap { transform(it).asSequence() }


fun <T : Any> Sequence<Any>.match(vararg expectedTypes: KClass<*>, last: KClass<T>): T? =
    (expectedTypes.asSequence() + last).zip(this + sequenceOf(null).cycle())
        .map { (expectedType, parent) -> parent?.takeIf(expectedType::isInstance) }
        .takeWhileInclusive { it != null }
        .lastOrNull()
        ?.let(last::cast)

private fun <T> Sequence<T>.cycle(): Sequence<T> = sequence { while (true) yieldAll(this@cycle) }

fun <T> Sequence<T>.takeWhileInclusive(predicate: (T) -> Boolean): Sequence<T> =
    sequence {
        for (elem in this@takeWhileInclusive) {
            yield(elem)
            if (!predicate(elem)) break
        }
    }

fun <T> sequenceOfLazyValues(vararg elements: () -> T): Sequence<T> = elements.asSequence().map { it() }